<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Ax WebLLM Chat</title>
    
    <!-- Load Tailwind CSS -->
    <script src="https://cdn.jsdelivr.net/npm/@tailwindcss/browser@4"></script>
    
    <!-- Load Ax library -->
    <!-- <script src="../ax/dist/index.global.js"></script> -->
 
    <script src="https://unpkg.com/@ax-llm/ax@latest?conditions=browser"></script>
    
    <!-- Load WebLLM -->
    <script type="module">
        import { MLCEngine } from "https://esm.run/@mlc-ai/web-llm";
        window.MLCEngine = MLCEngine;
        window.webllmLoaded = true;
    </script>
    
    <style>
        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap');
        body { font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif; }
        
        @keyframes pulse {
            0% { opacity: 1; }
            50% { opacity: 0.3; }
            100% { opacity: 1; }
        }
        .loading { animation: pulse 2s infinite; }
    </style>
</head>
<body class="bg-gray-50 min-h-screen p-4">
    <div class="max-w-2xl mx-auto">
        <!-- Header -->
        <div class="text-center mb-6">
            <h1 class="text-2xl font-semibold text-gray-900 mb-2">Ax WebLLM Chat</h1>
            <p class="text-gray-500 text-sm">Local AI inference in the browser</p>
        </div>

        <!-- Model Selection and Status -->
        <div class="bg-white rounded-lg shadow-sm border p-4 mb-6">
            <div class="mb-4">
                <label for="modelSelect" class="block text-sm font-medium text-gray-700 mb-2">Select Model:</label>
                <select id="modelSelect" class="w-full px-3 py-2 border rounded-lg text-sm focus:outline-none focus:ring-2 focus:ring-blue-500">
                    <option value="Llama-3.2-3B-Instruct-q4f32_1-MLC">Llama 3.2 3B Instruct (Default - Fast)</option>
                    <option value="Llama-3.2-1B-Instruct-q4f32_1-MLC">Llama 3.2 1B Instruct (Fastest)</option>
                    <option value="Llama-3.1-8B-Instruct-q4f32_1-MLC">Llama 3.1 8B Instruct (Better Quality)</option>
                    <option value="Llama-3.1-70B-Instruct-q4f16_1-MLC">Llama 3.1 70B Instruct (Best Quality - Slow)</option>
                    <option value="Mistral-7B-Instruct-v0.3-q4f32_1-MLC">Mistral 7B Instruct</option>
                    <option value="Phi-3.5-mini-instruct-q4f32_1-MLC">Phi 3.5 Mini Instruct</option>
                    <option value="gemma-2-2b-it-q4f32_1-MLC">Gemma 2 2B Instruct</option>
                    <option value="gemma-2-9b-it-q4f32_1-MLC">Gemma 2 9B Instruct</option>
                    <option value="Qwen2.5-0.5B-Instruct-q4f32_1-MLC">Qwen 2.5 0.5B Instruct</option>
                    <option value="Qwen2.5-1.5B-Instruct-q4f32_1-MLC">Qwen 2.5 1.5B Instruct</option>
                    <option value="Qwen2.5-3B-Instruct-q4f32_1-MLC">Qwen 2.5 3B Instruct</option>
                    <option value="Qwen2.5-7B-Instruct-q4f32_1-MLC">Qwen 2.5 7B Instruct</option>
                </select>
            </div>
            <div class="flex items-center gap-3">
                <div id="statusIndicator" class="w-3 h-3 rounded-full bg-red-400"></div>
                <span id="statusText">Initializing...</span>
                <button id="initBtn" class="ml-auto px-4 py-2 bg-blue-500 text-white rounded text-sm hover:bg-blue-600 disabled:bg-gray-300">
                    Load Model
                </button>
            </div>
            <div id="progressContainer" class="mt-3 hidden">
                <div class="w-full bg-gray-200 rounded-full h-2">
                    <div id="progressBar" class="bg-blue-500 h-2 rounded-full transition-all duration-300" style="width: 0%"></div>
                </div>
                <div id="progressText" class="text-xs text-gray-500 mt-1"></div>
            </div>
        </div>

        <!-- Chat -->
        <div class="bg-white rounded-lg shadow-sm border overflow-hidden">
            <!-- Messages -->
            <div id="messages" class="h-96 overflow-y-auto p-4 space-y-3">
                <div class="text-center text-gray-500 text-sm">
                    Load the model to start chatting
                </div>
            </div>
            
            <!-- Input -->
            <div class="border-t p-4">
                <div class="flex gap-3">
                    <input type="text" id="chatInput" placeholder="Type a message..." disabled 
                           class="flex-1 px-3 py-2 border rounded-lg text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 disabled:bg-gray-100">
                    <button id="sendBtn" disabled 
                            class="px-4 py-2 bg-blue-500 text-white rounded-lg text-sm hover:bg-blue-600 disabled:bg-gray-300">
                        Send
                    </button>
                </div>
            </div>
        </div>
    </div>

    <script>
        let ai = null;
        let chatBot = null;
        let engine = null;
        
        const elements = {
            statusIndicator: document.getElementById('statusIndicator'),
            statusText: document.getElementById('statusText'),
            initBtn: document.getElementById('initBtn'),
            modelSelect: document.getElementById('modelSelect'),
            progressContainer: document.getElementById('progressContainer'),
            progressBar: document.getElementById('progressBar'),
            progressText: document.getElementById('progressText'),
            messages: document.getElementById('messages'),
            chatInput: document.getElementById('chatInput'),
            sendBtn: document.getElementById('sendBtn')
        };
        
        function addMessage(content, type) {
            const div = document.createElement('div');
            div.className = `flex ${type === 'user' ? 'justify-end' : 'justify-start'}`;
            
            const messageClass = type === 'user' 
                ? 'bg-blue-500 text-white ml-12' 
                : 'bg-gray-100 text-gray-900 mr-12';
                
            div.innerHTML = `
                <div class="max-w-sm px-3 py-2 rounded-lg text-sm ${messageClass}">
                    ${content}
                </div>
            `;
            
            elements.messages.appendChild(div);
            elements.messages.scrollTop = elements.messages.scrollHeight;
        }
        
        function updateStatus(status, text) {
            const colors = {
                'loading': 'bg-yellow-400 loading',
                'ready': 'bg-green-400',
                'error': 'bg-red-400'
            };
            
            elements.statusIndicator.className = `w-3 h-3 rounded-full ${colors[status] || 'bg-gray-400'}`;
            elements.statusText.textContent = text;
        }
        
        async function initializeWebLLM() {
            try {
                elements.initBtn.disabled = true;
                elements.modelSelect.disabled = true;
                elements.progressContainer.classList.remove('hidden');
                
                const selectedModel = elements.modelSelect.value;
                updateStatus('loading', `Loading ${selectedModel}...`);
                
                // Initialize WebLLM engine
                engine = new window.MLCEngine();
                
                // Set up progress callback
                engine.setInitProgressCallback((progress) => {
                    const percentage = Math.round(progress.progress * 100);
                    elements.progressBar.style.width = `${percentage}%`;
                    elements.progressText.textContent = `${progress.text} (${percentage}%)`;
                });
                
                // Load the selected model
                await engine.reload(selectedModel);
                console.log('Model loaded, engine methods:', Object.getOwnPropertyNames(engine));
                console.log('Engine chat methods:', Object.getOwnPropertyNames(engine.chat || {}));
                console.log('Engine completions methods:', Object.getOwnPropertyNames(engine.chat?.completions || {}));
                
                // Initialize Ax AI with WebLLM
                console.log('Creating AxAI instance with engine:', engine);
                ai = new window.ax.AxAI({
                    name: 'webllm',
                    engine: engine,
                    config: {
                        model: selectedModel,
                        stream: false
                    }
                });
                console.log('AxAI instance created:', ai);
                
                // Create chat generator
                chatBot = window.ax.ax`
                    userMessage:${window.ax.f.string('User message')} -> 
                    assistantReply:${window.ax.f.string('Assistant response')}
                `;
                
                updateStatus('ready', 'Model loaded and ready!');
                elements.chatInput.disabled = false;
                elements.sendBtn.disabled = false;
                elements.progressContainer.classList.add('hidden');
                elements.initBtn.textContent = 'Ready';
                
                // Clear messages and add welcome
                elements.messages.innerHTML = '';
                addMessage('Hello! I\'m running locally in your browser. How can I help you?', 'assistant');
                
            } catch (error) {
                console.error('Failed to initialize WebLLM:', error);
                updateStatus('error', 'Failed to load model');
                elements.initBtn.disabled = false;
                elements.modelSelect.disabled = false;
                elements.progressContainer.classList.add('hidden');
            }
        }
        
        async function sendMessage() {
            const message = elements.chatInput.value.trim();
            if (!message || !ai || !chatBot) return;
            
            // Add user message
            addMessage(message, 'user');
            elements.chatInput.value = '';
            elements.sendBtn.disabled = true;
            elements.sendBtn.textContent = 'Thinking...';
            
            try {
                const response = await chatBot.forward(ai, {
                    userMessage: message
                });
                
                addMessage(response.assistantReply, 'assistant');
                
            } catch (error) {
                console.error('Chat error:', error);
                console.error('Full error details:', error);
                if (error.cause) {
                    console.error('Error cause:', error.cause);
                }
                addMessage(`Error: ${error.message}`, 'error');
            } finally {
                elements.sendBtn.disabled = false;
                elements.sendBtn.textContent = 'Send';
                elements.chatInput.focus();
            }
        }
        
        // Event listeners
        elements.initBtn.addEventListener('click', initializeWebLLM);
        elements.sendBtn.addEventListener('click', sendMessage);
        elements.chatInput.addEventListener('keypress', (e) => {
            if (e.key === 'Enter' && !elements.sendBtn.disabled) {
                sendMessage();
            }
        });
        
        // Initialize - wait for WebLLM to load
        setTimeout(() => {
            if (window.webllmLoaded && window.MLCEngine) {
                updateStatus('ready', 'Ready to load WebLLM model');
            } else {
                updateStatus('error', 'WebLLM failed to load');
            }
        }, 1000);
    </script>
</body>
</html>